package com.sean.learnelasticsearch.controller;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.SortOrder;
import co.elastic.clients.elasticsearch._types.aggregations.*;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.util.NamedValue;
import com.sean.learnelasticsearch.entity.Article;
import com.sean.learnelasticsearch.mapper.ArticleEsRepository;
import com.sean.learnelasticsearch.mapper.ArticleMapper;
import javafx.util.Pair;
import org.elasticsearch.search.aggregations.BucketOrder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchAggregation;
import org.springframework.data.elasticsearch.client.elc.NativeQuery;
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.query.Query;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;
import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 描述
 *
 * @author 林思瀚 (linsihan@cndatacom.com)
 * @date 2023-09-21 11:09:20
 */
@RestController
@RequestMapping("/es")
public class ESController {

    @Resource
    private ArticleEsRepository articleEsRepository;

    @Resource
    private ArticleMapper articleMapper;

    @Resource
    private ElasticsearchOperations operations;

    @Resource
    private ElasticsearchClient elasticsearchClient;

    /**
     * 分桶聚合
     *
     * @return List<Pair < String, Long>> 分桶结果
     * @throws IOException io 异常
     */
    @RequestMapping("/test01")
    public List<Pair<String, Long>> test01() throws IOException {
        List<Pair<String, Long>> list = new ArrayList<>();

        List<NamedValue<SortOrder>> sortOrderList = new ArrayList<>();

        // 添加排序规则，可以多次调用 add 方法添加多个规则
        sortOrderList.add(new NamedValue<>("_count", SortOrder.Desc));
        Query query = NativeQuery.builder()
                .withMaxResults(0)
                .withAggregation("agg_tag", Aggregation.of(a -> a
                        .terms(ta -> ta.field("type").size(5).order(sortOrderList))
                ))
                .build();
        // 查找结果集
        SearchHits<Article> searchHits = operations.search(query, Article.class);
        // 聚合结果集
        List<ElasticsearchAggregation> aggregations = (List<ElasticsearchAggregation>) searchHits.getAggregations().aggregations();
        Aggregate aggregate = aggregations.get(0).aggregation().getAggregate();
        StringTermsAggregate stringTermsAggregate = (StringTermsAggregate) aggregate._get();
        // 桶集合
        List<StringTermsBucket> buckets = (List<StringTermsBucket>) stringTermsAggregate.buckets()._get();
        // 遍历桶
        for (StringTermsBucket bucket : buckets) {
            list.add(new Pair<>(bucket.key().stringValue(), bucket.docCount()));
        }
        return list;
    }

    /**
     * CRUD (Read)
     *
     * @return Article 文章 toString
     */
    @GetMapping("/test02")
    public Article testFindById() {
        return articleEsRepository.findById(1705952826721891065L).get();
    }

    /**
     * 更好的分桶聚合（使用 ES 客户端）
     *
     * @return List<Pair < String, Long>> 分桶结果
     */
    @GetMapping("/test03")
    public List<Pair<String, Long>> test03() {

        List<Pair<String, Long>> list = new ArrayList<>();
        List<NamedValue<SortOrder>> sortOrderList = new ArrayList<>();

        // 添加排序规则，可以添加多个排序规则
        sortOrderList.add(new NamedValue<>("_count", SortOrder.Desc));
        SearchResponse<Void> response = null;
        try {
            response = elasticsearchClient.search(s -> s
                            .index("article_index")     // 索引
                            .size(0)    // 不显示 hits 数据：size:0
                            .aggregations("agg_tag", a -> a // 聚合（名称，查询体）
                                    .terms(t -> t   // 分桶聚合(查询体)
                                            .field("type")  // 分桶字段
                                            .size(5)    // 桶的数量
                                            .order(sortOrderList)   // 排序规则
                                    )
                            ),
                    Void.class
            );
            // 得到桶列表
            List<StringTermsBucket> termsBucketList = response.aggregations().get("agg_tag")
                    .sterms().buckets().array();
            // 遍历桶
            for (StringTermsBucket bucket : termsBucketList) {
                list.add(new Pair<>(bucket.key().stringValue(), bucket.docCount()));
            }
            return list;
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @GetMapping("searchMatch/{searchWord}/{pageNo}/{pageSize}")
    public List<Article> searchMatch(@PathVariable("searchWord") String searchWord,
                                     @PathVariable("pageNo") Integer pageNo,
                                     @PathVariable("pageSize") Integer pageSize) {
        List<Article> content = articleEsRepository
                .findByTitleOrContext(searchWord, searchWord, PageRequest.of(pageNo, pageSize)).getContent();
        return content;
    }

    @GetMapping("searchMatchClient/{searchWord}/{pageNo}/{pageSize}")
    public List<Article> searchMatchClient(@PathVariable("searchWord") String searchWord,
                                           @PathVariable("pageNo") Integer pageNo,
                                           @PathVariable("pageSize") Integer pageSize) {
        SearchResponse<Article> response = null;
        try {
            response = elasticsearchClient.search(s -> s
                    .index("article_index")
                    .query(q -> q
                            .match(t -> t
                                    .field("title")
                                    .query(searchWord)
                            )
                    )
                    .from(2)
                    .size(3), Article.class
            );
        } catch (IOException e) {
            e.printStackTrace();
        }
        List<Article> articles = new ArrayList<>();
        assert response != null;
        for (Hit<Article> hit : response.hits().hits()) {
            articles.add(hit.source());
        }
        return articles;
    }

    @GetMapping("multiMatch/{searchWord}/{pageNo}/{pageSize}")
    public List<Article> multiMatch(@PathVariable("searchWord") String searchWord,
                                    @PathVariable("pageNo") Integer pageNo,
                                    @PathVariable("pageSize") Integer pageSize) {
        List<String> fields = new ArrayList<>();
        fields.add("title");
        fields.add("context");
        SearchResponse<Article> response = null;
        try {
            response = elasticsearchClient.search(s -> s
                    .index("article_index")
                    .query(q -> q
                            .multiMatch(t -> t
                                    .fields(fields)
                                    .query(searchWord)
                            )
                    )
                    .from(2)
                    .size(3), Article.class
            );
        } catch (IOException e) {
            e.printStackTrace();
        }
        List<Article> articles = new ArrayList<>();
        assert response != null;
        for (Hit<Article> hit : response.hits().hits()) {
            articles.add(hit.source());
        }
        return articles;
    }

    @GetMapping("/avg")
    public List<Pair<String, Double>> avg() {
        SearchResponse<Article> response = null;
        try {
            response = elasticsearchClient.search(s -> s
                    .index("article_index")
                    .size(0)
                    .aggregations("max_like_count", a -> a
                            .max(t -> t
                                    .field("likeCount")
                            )
                    )
                    .aggregations("min_like_count", a -> a
                            .min(t -> t
                                    .field("likeCount")
                            )
                    )
                    .aggregations("avg_like_count", a -> a
                            .avg(t -> t
                                    .field("likeCount")
                            )
                    ), Article.class);
        } catch (IOException e) {
            e.printStackTrace();
        }
        // 封装指标聚合结果集
        List<Pair<String, Double>> list = new ArrayList<>();
        list.add(new Pair<String, Double>("max", response.aggregations().get("max_like_count").max().value()));
        list.add(new Pair<String, Double>("min", response.aggregations().get("min_like_count").min().value()));
        list.add(new Pair<String, Double>("avg", response.aggregations().get("avg_like_count").avg().value()));
        return list;
    }

    @GetMapping("/pineline")
    public Pair<String, Double> pineline() {
        SearchResponse<Article> response = null;
        try {
            response = elasticsearchClient.search(s -> s
                            .index("article_index")
                            .size(0)
                            .aggregations("type_agg", a->a
                                    .terms(t->t
                                            .field("type")
                                    )
                                    .aggregations("avg_like_count", agg->agg
                                            .avg(t->t
                                                    .field("likeCount")
                                            )
                                    )
                            )
                            .aggregations("min_bucket", a->a
                                    .minBucket(b->b
                                            .bucketsPath(p->p.single("type_agg>avg_like_count"))
                                    )
                            )
                            , Article.class);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return new Pair<String, Double>("min_bucket", response.aggregations().get("min_bucket").bucketMetricValue().value());
    }
}
